import sqlite3
from sqlite3 import OperationalError
import pandas as pd
from sqlalchemy import create_engine

# 使用条件导入避免模块缺失导致的错误
try:
    import mysql.connector
except ImportError:
    mysql = None

try:
    import psycopg2
except ImportError:
    psycopg2 = None

try:
    import pymssql
except ImportError:
    pymssql = None

class DatabaseConnector:
    """数据库连接工具类，支持多种数据库类型"""
    
    def __init__(self, db_type, host=None, port=None, database=None, username=None, password=None):
        self.db_type = db_type
        self.host = host
        self.port = port
        self.database = database
        self.username = username
        self.password = password
        self.connection = None
    
    def get_connection(self):
        """获取数据库连接"""
        if self.connection:
            return self.connection
            
        try:
            if self.db_type == 'mysql':
                if 'mysql' not in globals() or mysql is None:
                    raise ImportError("mysql-connector-python 模块未安装")
                self.connection = mysql.connector.connect(
                    host=self.host,
                    port=self.port,
                    database=self.database,
                    user=self.username,
                    password=self.password
                )
            elif self.db_type == 'postgresql':
                if psycopg2 is None:
                    raise ImportError("psycopg2 模块未安装")
                self.connection = psycopg2.connect(
                    host=self.host,
                    port=self.port,
                    dbname=self.database,
                    user=self.username,
                    password=self.password
                )
            elif self.db_type == 'sqlserver':
                if pymssql is None:
                    raise ImportError("pymssql 模块未安装")
                conn_str = f"mssql+pymssql://{self.username}:{self.password}@{self.host}:{self.port}/{self.database}"
                engine = create_engine(conn_str)
                self.connection = engine.connect()
            elif self.db_type == 'sqlite':
                self.connection = sqlite3.connect(self.database)
            else:
                raise ValueError(f"不支持的数据库类型: {self.db_type}")
                
            return self.connection
        except Exception as e:
            print(f"数据库连接失败: {str(e)}")
            return None
    
    def test_connection(self):
        """测试数据库连接是否有效"""
        try:
            conn = self.get_connection()
            if conn:
                conn.close()
                self.connection = None  # 测试连接后关闭
                return True
            return False
        except:
            return False
    
    def execute_query(self, query):
        """执行查询并返回结果"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            # 执行查询
            cursor.execute(query)
            
            # 获取列名
            columns = [desc[0] for desc in cursor.description] if cursor.description else []
            
            # 获取结果
            result = cursor.fetchall()
            
            # 对于写操作，需要提交事务
            if query.strip().upper().startswith(('INSERT', 'UPDATE', 'DELETE', 'CREATE', 'ALTER', 'DROP')):
                conn.commit()
                
            return result, columns
        except Exception as e:
            # 发生错误时回滚
            if self.connection:
                self.connection.rollback()
            raise Exception(f"查询执行失败: {str(e)}")
        finally:
            # 关闭游标
            if 'cursor' in locals():
                cursor.close()
    
    def execute_script(self, script):
        """执行SQL脚本（多个语句）"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            cursor.executescript(script)
            conn.commit()
            return True
        except Exception as e:
            if self.connection:
                self.connection.rollback()
            raise Exception(f"脚本执行失败: {str(e)}")
        finally:
            if 'cursor' in locals():
                cursor.close()
    
    def to_dataframe(self, query):
        """执行查询并返回Pandas DataFrame"""
        try:
            conn = self.get_connection()
            df = pd.read_sql(query, conn)
            return df
        except Exception as e:
            raise Exception(f"转换为DataFrame失败: {str(e)}")
    
    def close(self):
        """关闭数据库连接"""
        if self.connection:
            self.connection.close()
            self.connection = None
    
    def __del__(self):
        """对象销毁时关闭连接"""
        self.close()
    